// Copyright 2017 Baidu Inc. All Rights Reserved.
// Author: Su Hao (suhao@baidu.com)
//
// Description: The encrypted network I/O.
//

#include "baidu_ca_trans_encrypted.h"

#if !defined(MBEDTLS_CONFIG_FILE)
#include "mbedtls/config.h"
#else
#include MBEDTLS_CONFIG_FILE
#endif
#include "mbedtls/ssl.h"
#include "mbedtls/entropy.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/timing.h"
#include "mbedtls/x509_crt.h"
#include "mbedtls/net.h"
#include "mbedtls/debug.h"

#include "baidu_ca_internal.h"
#include "baidu_ca_trans_wrapper.h"
#include "baidu_ca_timestamp.h"
#include "baidu_ca_random.h"
#include "duer_log.h"


#define DRBG_PERS_CLIENT            "BCA DTLS CLIENT"
#define COMMON_NAME                 "*.iot.baidu.com"

#define HANDSHAKE_TIMEOUT_MIN       (1000)
#define HANDSHAKE_TIMEOUT_MAX       (4000)

// Suppress Compiler warning Function declared never referenced
#define SUPPRESS_WARNING(func) (void)(func)

typedef enum _baidu_ca_trans_status_enum {
    BCA_TRANS_ST_DISCONNECTED,
    BCA_TRANS_ST_CONNECTED,
    BCA_TRANS_ST_HANDSHAKED,
} bca_trans_status_e;

typedef struct _bca_trans_encrypted_s {
#ifndef MBEDTLS_DRBG_ALT
    mbedtls_entropy_context         entropy;
    mbedtls_ctr_drbg_context        ctr_drbg;
#endif
    mbedtls_x509_crt                cacert;
    mbedtls_ssl_context             ssl;
    mbedtls_ssl_config              ssl_conf;
    mbedtls_timing_delay_context    timer;

    bca_u8_t                        status;
} bca_trans_encrypted_t, *bca_trans_encrypted_ptr;

BCA_LOC_IMPL bca_status_t bca_trans_mbedtls2status(int status) {
    bca_status_t rs = status;

    if (status == MBEDTLS_ERR_SSL_TIMEOUT) {
        rs = BCA_ERR_TRANS_TIMEOUT;
    } else if (status == MBEDTLS_ERR_SSL_WANT_READ
               || status == MBEDTLS_ERR_SSL_WANT_WRITE) {
        rs = BCA_ERR_TRANS_WOULD_BLOCK;
    } else if (status < 0) {
        rs = BCA_ERR_TRANS_INTERNAL_ERROR;
    } else {
        // do nothing
    }

    if (rs < 0 && rs != BCA_ERR_TRANS_WOULD_BLOCK) {
        DUER_LOGE("mbedtls error: %d(-0x%04x)", status, -status);
    }

    return rs;
}

#if defined(BCA_MBEDTLS_DEBUG) && (BCA_MBEDTLS_DEBUG > 0)
BCA_LOC_IMPL void bca_trans_encrypted_wrap_debug(void* ctx, int level,
        const char* file, int line, const char* msg) {
    bca_debug_print(level, file, line, msg);
}
#endif

BCA_LOC_IMPL int bca_trans_encrypted_wrap_send(void* ctx,
        const unsigned char* buf, size_t len) {
    bca_status_t rs = bca_trans_wrapper_send((bca_trans_ptr)ctx, buf, len, NULL);

    if (rs == BCA_ERR_TRANS_WOULD_BLOCK) {
        return MBEDTLS_ERR_SSL_WANT_WRITE;
    } else if (rs < BCA_NO_ERR) {
        return MBEDTLS_ERR_NET_SEND_FAILED;
    } else {
        // do nothing
    }

    return rs;
}

BCA_LOC_IMPL int bca_trans_encrypted_wrap_recv(void* ctx, unsigned char* buf,
        size_t len) {
    bca_status_t rs = BCA_ERR_INTERNAL;
    bca_trans_ptr trans = (bca_trans_ptr)ctx;
    rs = bca_trans_wrapper_recv(trans, buf, len, NULL);

    if (rs == BCA_ERR_TRANS_WOULD_BLOCK) {
        return MBEDTLS_ERR_SSL_WANT_READ;
    } else if (rs < BCA_NO_ERR) {
        return MBEDTLS_ERR_NET_RECV_FAILED;
    } else {
        // do nothing
    }

    return rs;
}

BCA_LOC_IMPL int bca_trans_encrypted_wrap_recv_timeout(void* ctx,
        unsigned char* buf, size_t len, uint32_t timeout) {
    bca_status_t rs = BCA_ERR_INTERNAL;
    bca_trans_ptr trans = (bca_trans_ptr)ctx;
    DUER_LOGV("bca_trans_encrypted_wrap_recv_timeout: timeout = %d", timeout);
    rs = bca_trans_wrapper_recv_timeout(trans, buf, len, timeout, NULL);

    if (rs == BCA_ERR_TRANS_WOULD_BLOCK) {
        return MBEDTLS_ERR_SSL_WANT_READ;
    } else if (rs == BCA_ERR_TRANS_TIMEOUT) {
        return MBEDTLS_ERR_SSL_TIMEOUT;
    } else if (rs < BCA_NO_ERR) {
        return MBEDTLS_ERR_NET_RECV_FAILED;
    } else {
        // do nothing
    }

    return rs;
}

#ifndef MBEDTLS_DRBG_ALT
BCA_LOC_IMPL int bca_trans_encrypted_poll(void* ctx, unsigned char* output,
        size_t len, size_t* olen) {
    for (uint16_t i = 0; i < len; i++) {
        output[i] = bca_random() & 0xFF;
    }

    *olen = len;
    DUER_LOGD("bca_trans_encrypted_poll: len = %d", len);
    return 0;
}
#endif

BCA_LOC_IMPL unsigned long bca_trans_encrypted_get_timer(
        struct mbedtls_timing_hr_time* val, int reset) {
    unsigned long delta;
    bca_u32_t offset = bca_timestamp();
    bca_u32_t* t = (bca_u32_t*)val;

    if (reset) {
        if (t) {
            *t  = offset;
        }

        return 0;
    }

    delta = offset - *t;
    return delta;
}

BCA_LOC_IMPL void bca_trans_encrypted_set_delay(void* data, uint32_t int_ms,
        uint32_t fin_ms) {
    mbedtls_timing_delay_context* ctx = (mbedtls_timing_delay_context*) data;
    ctx->int_ms = int_ms;
    ctx->fin_ms = fin_ms;

    if (fin_ms != 0) {
        bca_trans_encrypted_get_timer(&ctx->timer, 1);
    }
}

BCA_LOC_IMPL int bca_trans_encrypted_get_delay(void* data) {
    mbedtls_timing_delay_context* ctx = (mbedtls_timing_delay_context*)data;
    unsigned long elapsed_ms;

    if (ctx->fin_ms == 0) {
        return -1;
    }

    elapsed_ms = bca_trans_encrypted_get_timer(&ctx->timer, 0);

    if (elapsed_ms >= ctx->fin_ms) {
        return 2;
    }

    if (elapsed_ms >= ctx->int_ms) {
        return 1;
    }

    return 0;
}

BCA_LOC_IMPL bca_trans_encrypted_ptr bca_trans_get_context(bca_trans_ptr trans) {
    bca_trans_encrypted_ptr ptr = NULL;
    int rs = 0;
    bca_bool use_read_timeout = 0;
    DUER_LOGV("==> bca_trans_get_context: trans = %p", trans);

    if (!trans) {
        goto exit;
    }

    ptr = (bca_trans_encrypted_ptr)trans->secure;

    if (ptr) {
        goto exit;
    }

    ptr = (bca_trans_encrypted_ptr)BCA_MALLOC(sizeof(bca_trans_encrypted_t));

    if (!ptr) {
        goto exit;
    }

    BCA_MEMSET(ptr, 0, sizeof(bca_trans_encrypted_t));
    /*
     * 0. Initialize the RNG and the session data
     */
    mbedtls_ssl_init(&ptr->ssl);
    mbedtls_ssl_config_init(&ptr->ssl_conf);
    mbedtls_x509_crt_init(&ptr->cacert);
#ifndef MBEDTLS_DRBG_ALT
    mbedtls_ctr_drbg_init(&ptr->ctr_drbg);
    mbedtls_entropy_init(&ptr->entropy);
    SUPPRESS_WARNING(bca_trans_encrypted_poll);
    rs = mbedtls_entropy_add_source(&ptr->entropy, bca_trans_encrypted_poll, NULL, 128, 1);

    if (rs != 0) {
        DUER_LOGE("mbedtls_entropy_add_source failed: %d", rs);
        goto error;
    }

    rs = mbedtls_ctr_drbg_seed(&ptr->ctr_drbg,
                               mbedtls_entropy_func,
                               &ptr->entropy,
                               (const unsigned char*)(DRBG_PERS_CLIENT),
                               sizeof(DRBG_PERS_CLIENT));

    if (rs != 0) {
        DUER_LOGE("mbedtls_ctr_drbg_seed failed: %d", rs);
        goto error;
    }

#endif
    /*
     * 1. Load certificates
     */
    rs = mbedtls_x509_crt_parse(&ptr->cacert,
                                (const unsigned char*)trans->cert,
                                trans->cert_len);

    if (rs != 0) {
        DUER_LOGE("mbedtls_x509_crt_parse failed: %d", rs);
        goto error;
    }

    /*
     * 2. Setup stuff
     */
    rs = mbedtls_ssl_config_defaults(&ptr->ssl_conf,
                                     MBEDTLS_SSL_IS_CLIENT,
                                     trans->addr.type == BCA_PROTO_UDP
                                     ? MBEDTLS_SSL_TRANSPORT_DATAGRAM
                                     : MBEDTLS_SSL_TRANSPORT_STREAM,
                                     MBEDTLS_SSL_PRESET_DEFAULT);

    if (rs != 0) {
        DUER_LOGE("mbedtls_ssl_config_defaults failed: %d", rs);
        goto error;
    }

    /* OPTIONAL is usually a bad choice for security, but makes interop easier
     * in this simplified example, in which the ca chain is hardcoded.
     * Production code should set a proper ca chain and use REQUIRED. */
    mbedtls_ssl_conf_authmode(&ptr->ssl_conf, MBEDTLS_SSL_VERIFY_OPTIONAL);
    //mbedtls_ssl_conf_authmode(&ssl_conf, MBEDTLS_SSL_VERIFY_REQUIRED);
    mbedtls_ssl_conf_ca_chain(&ptr->ssl_conf, &ptr->cacert, NULL);
#ifndef MBEDTLS_DRBG_ALT
    mbedtls_ssl_conf_rng(&ptr->ssl_conf, mbedtls_ctr_drbg_random, &ptr->ctr_drbg);
#else
    mbedtls_ssl_conf_rng(&ptr->ssl_conf, mbedtls_ctr_drbg_random, NULL);
#endif
    //config-option it by user
    //mbedtls_ssl_conf_read_timeout(&_ssl_conf, REPORT_NET_RECEIVE_TIMEOUT);
#if defined(BCA_MBEDTLS_DEBUG) && (BCA_MBEDTLS_DEBUG > 0)
    SUPPRESS_WARNING(bca_trans_encrypted_wrap_debug);
    mbedtls_ssl_conf_dbg(&ptr->ssl_conf, bca_trans_encrypted_wrap_debug, NULL);
    mbedtls_debug_set_threshold(BCA_MBEDTLS_DEBUG);
#endif
    rs = mbedtls_ssl_setup(&ptr->ssl, &ptr->ssl_conf);

    if (rs != 0) {
        DUER_LOGE("mbedtls_ssl_setup failed: %d", rs);
        goto error;
    }

    rs = mbedtls_ssl_set_hostname(&ptr->ssl, COMMON_NAME);

    if (rs != 0) {
        DUER_LOGE("mbedtls_ssl_setup failed: %d", rs);
        goto error;
    }

    use_read_timeout = trans->read_timeout != BCA_READ_FOREVER;
    mbedtls_ssl_set_bio(&ptr->ssl, trans,
                        bca_trans_encrypted_wrap_send, bca_trans_encrypted_wrap_recv,
                        use_read_timeout ? bca_trans_encrypted_wrap_recv_timeout : NULL);
    mbedtls_ssl_set_timer_cb(&ptr->ssl,
                             &ptr->timer,
                             bca_trans_encrypted_set_delay,
                             bca_trans_encrypted_get_delay);

    if (use_read_timeout) {
        mbedtls_ssl_conf_read_timeout(&ptr->ssl_conf, trans->read_timeout);
    }
#if defined(MBEDTLS_SSL_PROTO_DTLS)
    mbedtls_ssl_conf_handshake_timeout(&ptr->ssl_conf, HANDSHAKE_TIMEOUT_MIN,
                                       HANDSHAKE_TIMEOUT_MAX);
#endif
    ptr->status = BCA_TRANS_ST_DISCONNECTED;
    trans->secure = ptr;
exit:
    DUER_LOGV("<== bca_trans_get_context: ptr = %p", ptr);
    return ptr;
error:

    if (ptr) {
        BCA_FREE(ptr);
        ptr = NULL;
    }

    goto exit;
}

BCA_LOC_IMPL bca_status_t bca_trans_encrypted_handshake(bca_trans_ptr trans) {
    bca_status_t rs = BCA_ERR_INTERNAL;
    bca_trans_encrypted_ptr ptr = trans ? (bca_trans_encrypted_ptr)trans->secure : NULL;
    DUER_LOGV("==> bca_trans_encrypted_handshake: trans = %p", trans);

    if (!ptr) {
        goto exit;
    }

    rs = mbedtls_ssl_handshake(&ptr->ssl);
    DUER_LOGD("    bca_trans_encrypted_handshake: rs = %d(-0x%04x)", rs, -rs);
    rs = bca_trans_mbedtls2status(rs);

    if (rs == BCA_NO_ERR) {
        ptr->status = BCA_TRANS_ST_HANDSHAKED;
    }

exit:
    DUER_LOGV("<== bca_trans_encrypted_handshake: rs = %d", rs);
    return rs;
}

BCA_INT_IMPL bca_status_t bca_trans_encrypted_do_connect(bca_trans_ptr trans,
        const bca_addr_t* addr) {
    bca_status_t rs = BCA_ERR_INTERNAL;
    bca_trans_encrypted_ptr ptr = NULL;
    DUER_LOGV("==> bca_trans_encrypted_do_connect: trans = %p", trans);
    ptr = bca_trans_get_context(trans);

    if (!ptr) {
        goto exit;
    }

    rs = bca_trans_wrapper_connect(trans, addr);

    if (rs == BCA_NO_ERR) {
        ptr->status = BCA_TRANS_ST_CONNECTED;
    }

exit:
    DUER_LOGV("<== bca_trans_encrypted_do_connect: rs = %d", rs);
    return rs;
}

BCA_LOC_IMPL bca_status_t bca_trans_encrypted_do_prepare(bca_trans_ptr trans) {
    bca_status_t rs = BCA_ERR_INTERNAL;
    bca_trans_encrypted_ptr ptr = bca_trans_get_context(trans);
    DUER_LOGV("==> bca_trans_encrypted_do_prepare: trans = %p", trans);

    if (!ptr) {
        DUER_LOGW("    bca_trans_encrypted_do_prepare: ptr is null!!!");
        goto exit;
    }

    switch (ptr->status) {
    case BCA_TRANS_ST_DISCONNECTED:
        rs = bca_trans_encrypted_do_connect(trans, &trans->addr);
        break;

    case BCA_TRANS_ST_CONNECTED:
        rs = bca_trans_encrypted_handshake(trans);
        break;

    case BCA_TRANS_ST_HANDSHAKED:
        rs = BCA_NO_ERR;
        break;
    default:
        // do nothing
        break;
    }

exit:
    DUER_LOGV("<== bca_trans_encrypted_do_prepare: rs = %d", rs);
    return rs;
}

BCA_LOC_IMPL bca_status_t bca_trans_encrypted_prepare(bca_trans_ptr trans) {
    bca_status_t rs = BCA_ERR_INTERNAL;
    bca_trans_encrypted_ptr ptr = bca_trans_get_context(trans);
    DUER_LOGV("==> bca_trans_encrypted_prepare: trans = %p", trans);

    if (!ptr) {
        DUER_LOGW("    bca_trans_encrypted_prepare: ptr is null!!!");
        goto exit;
    }

    rs = BCA_NO_ERR;

    while (ptr->status != BCA_TRANS_ST_HANDSHAKED && rs == BCA_NO_ERR) {
        rs = bca_trans_encrypted_do_prepare(trans);
    }

exit:
    DUER_LOGV("<== bca_trans_encrypted_prepare: rs = %d", rs);
    return rs;
}

BCA_INT_IMPL bca_status_t bca_trans_encrypted_connect(bca_trans_ptr trans,
        const bca_addr_t* addr) {
    bca_status_t rs = BCA_ERR_INTERNAL;
    bca_trans_encrypted_ptr ptr = NULL;
    DUER_LOGV("==> bca_trans_encrypted_connect: trans = %p", trans);
    ptr = bca_trans_get_context(trans);

    if (!ptr) {
        goto exit;
    }

    if (ptr->status == BCA_TRANS_ST_DISCONNECTED) {
        rs = bca_trans_encrypted_do_connect(trans, addr);

        if (rs < BCA_NO_ERR) {
            goto exit;
        }
    }

    rs = bca_trans_encrypted_prepare(trans);
exit:
    DUER_LOGV("<== bca_trans_encrypted_connect: rs = %d", rs);
    return rs;
}

BCA_INT_IMPL bca_status_t bca_trans_encrypted_send(bca_trans_ptr trans,
                                                   const void* data,
                                                   bca_size_t size,
                                                   const bca_addr_t* addr) {
    bca_trans_encrypted_ptr ptr = trans ? (bca_trans_encrypted_ptr)trans->secure : NULL;
    int rs = BCA_ERR_INTERNAL;
    DUER_LOGV("==> bca_trans_encrypted_send: trans = %p, size:%d", trans, size);
    rs = bca_trans_encrypted_prepare(trans);

    if (rs < BCA_NO_ERR) {
        DUER_LOGW("    bca_trans_encrypted_send: prepare failed, because rs = %d", rs);
        goto exit;
    }

    rs = mbedtls_ssl_write(&(ptr->ssl), (const unsigned char*)data, size);
    rs = bca_trans_mbedtls2status(rs);
    DUER_LOGV("<== bca_trans_encrypted_send: rs = %d", rs);
exit:
    return rs;
}

BCA_INT_IMPL bca_status_t bca_trans_encrypted_recv(bca_trans_ptr trans,
                                                   void* data,
                                                   bca_size_t size,
                                                   bca_addr_t* addr) {
    bca_trans_encrypted_ptr ptr = trans ? (bca_trans_encrypted_ptr)trans->secure : NULL;
    int rs = BCA_ERR_INTERNAL;
    DUER_LOGV("==> bca_trans_encrypted_recv: trans = %p", trans);
    rs = bca_trans_encrypted_prepare(trans);
    if (rs < BCA_NO_ERR) {
        DUER_LOGW("    bca_trans_encrypted_recv: prepare failed, because rs = %d", rs);
        goto exit;
    }

    rs = mbedtls_ssl_read(&(ptr->ssl), (unsigned char*)data, size);
    rs = bca_trans_mbedtls2status(rs);
    DUER_LOGV("<== bca_trans_encrypted_recv: rs = %d", rs);
exit:
    return rs;
}

BCA_INT_IMPL bca_status_t bca_trans_encrypted_close(bca_trans_ptr trans) {
    bca_trans_encrypted_ptr ptr = trans ? (bca_trans_encrypted_ptr)trans->secure : NULL;
    DUER_LOGV("==> bca_trans_encrypted_close: trans = %p", trans);

    if (ptr) {
        mbedtls_ssl_close_notify(&ptr->ssl);
        mbedtls_ssl_session_reset(&ptr->ssl);
        bca_trans_wrapper_close(trans);
        mbedtls_x509_crt_free(&ptr->cacert);
        mbedtls_ssl_free(&ptr->ssl);
        mbedtls_ssl_config_free(&ptr->ssl_conf);
#ifndef MBEDTLS_DRBG_ALT
        mbedtls_ctr_drbg_free(&ptr->ctr_drbg);
        mbedtls_entropy_free(&ptr->entropy);
#endif
        BCA_FREE(ptr);
    }

    DUER_LOGV("<== bca_trans_encrypted_close");
    return BCA_NO_ERR;
}

BCA_INT_IMPL bca_status_t bca_trans_encrypted_set_read_timeout(
        bca_trans_ptr trans, bca_u32_t timeout) {
    bca_trans_encrypted_ptr ptr = trans ? (bca_trans_encrypted_ptr)trans->secure : NULL;

    if (ptr) {
        if (timeout != BCA_READ_FOREVER) {
            mbedtls_ssl_conf_read_timeout(&ptr->ssl_conf, timeout);
        }
        return BCA_NO_ERR;
    }

    return BCA_ERR_INTERNAL;
}
